package rate_limiter

import (
	"context"
	"google.golang.org/grpc"
)

var handlers = map[string]func(ctx context.Context, req interface{}) (any, error){}

type CustomRatelimitHandler struct {
	handlers map[string]func(ctx context.Context, req interface{})
}

func (h *CustomRatelimitHandler) UnaryServerInterceptor(tb *TokenBucket) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		// 这个代表的是没有触发限流
		if tb.Consume(1) {
			return handler(ctx, req)
		}
		// 触发了限流怎么样？
		// 这里你应该转异步
		// 你没有足够的细节，执行转异步
		// 我搞一个简单的标记位，我让业务方知道需要转就可以

		//ctx = context.WithValue(ctx, "RateLimit", true)
		strategy := handlers[info.FullMethod]
		return strategy(ctx, req)
	}
}

func UnaryServerInterceptor(tb *TokenBucket) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		// 这个代表的是没有触发限流
		if tb.Consume(1) {
			return handler(ctx, req)
		}
		// 触发了限流怎么样？
		// 这里你应该转异步
		// 你没有足够的细节，执行转异步
		// 我搞一个简单的标记位，我让业务方知道需要转就可以

		ctx = context.WithValue(ctx, "RateLimited", true)
		//strategy := handlers[info.FullMethod]
		//return strategy(ctx, req)
		return handler(ctx, req)
	}
}
